package com.feidee.fd.sml.algorithm.component.ml.ps.classification

import java.io.File

import com.feidee.fd.sml.algorithm.component.ml.MLParam
import com.tencent.angel.sona.ml.PipelineStage
import com.tencent.angel.sona.ml.classification.{AngelClassifier, CusAngelClassifier}
import com.tencent.angel.sona.ml.param.HasModelName
import ml.combust.mleap.core.classification.NaiveBayesModel.ModelType

/**
  * @author xiongjun
  * @date 2019/11/7 14:39
  * @description
  * @reviewer
  */
case class PSLRParam(
                          override val input_pt: String,
                          override val output_pt: String,
                          override val hive_table: String,
                          override val flow_time: String,
                          override val modelPath: String,
                          override var labelCol: String,
                          override val featuresCol: String,
                          override var predictionCol: String,
                          override val metrics: Array[String],
                          // 样本权重列
//                          weightCol: String,
                          rawPredictionCol: String,
                          probabilityCol: String,
                          //学习率,默认0.1
                          lr:Double,
                          // 最大迭代次数，默认 100
                          maxIter: Int,
                          isSparse:Boolean,
                          jsonFile:String= "./logreg.json",
                          modelName:String="LR"
                        ) extends MLParam {
  def this() = this(null, null, null, null, null, "label", "features", "prediction", Array.empty[String],/*null,*/"rawPrediction","probability",0.1,100,true,"./logreg.json","LR")

  override def verify(): Unit = {
    super.verify()
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("probabilityCol" -> probabilityCol)
    map += ("rawPredictionCol" -> rawPredictionCol)
    map += ("lr" -> lr)
    map += ("maxIter" -> maxIter)
    map += ("isSparse" -> isSparse)
    map += ("jsonFile" -> jsonFile)
    map += ("modelName" -> modelName)
    map
  }
}
class PSLRComponent extends AbstractPSClassificationComponent[PSLRParam] {
  override def setUp(param: PSLRParam): PipelineStage = {
    val classifier = new CusAngelClassifier()
      .setMaxIter(param.maxIter)
      .setLearningRate(param.lr)
      .setFeaturesCol(param.featuresCol)
      .setPredictionCol(param.predictionCol)
      .setRawPredictionCol(param.rawPredictionCol)
      .setLabelCol(param.labelCol)
      .setModelJsonFile(param.jsonFile)
    classifier
  }
}
object PSLRComponent {
  def apply(paramStr: String): Unit = {
    new PSLRComponent()(paramStr)
  }

  def main(args: Array[String]): Unit = {
    PSLRComponent(args(0))
  }
}
